package com.practice.security.aes;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.util.Base64;

import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;

public class AESKeyGenerator {

	private static final String SECRET_KEY_PATH = "./AES/secretKey.bin";

	private static final String INIT_VECTOR_PATH = "./AES/initVector.bin";

	private static final String SECRET_KEY_BASE64_PATH = "./AES/secretKey.txt";

	private static final String INIT_VECTOR_BASE64_PATH = "./AES/initVector.txt";

	private static final int SALT_LENGTH = 16;

	private static final String PASSWORD = "password";

	private static final int ITERATION_COUNT = 0x10000;

	// Key size can be 16 bytes, 24 bytes, 32 bytes.
	private static final int KEY_LENGTH = 256;

	private static final String PBKDF2_ALGORITHM = "PBKDF2WithHmacSHA1";

	private static final int INIT_VECTOR_LENGTH = 16;

	public static void main(String[] args) throws IOException {

		SecretKey secretKey = generateKey();
		if (secretKey == null) {
			return;
		}

		writeSecretKey(secretKey);
		
		byte[] initVectorBytes = generateInitVector();
		writeInitVector(initVectorBytes);
	}

	public static SecretKey generateKey() {

		SecureRandom random = new SecureRandom();
		byte[] salt = new byte[SALT_LENGTH];
		random.nextBytes(salt);

		SecretKey secretKey = null;
		try {
			KeySpec spec = new PBEKeySpec(PASSWORD.toCharArray(), salt, ITERATION_COUNT, KEY_LENGTH);
			SecretKeyFactory secretKeyFactory = SecretKeyFactory.getInstance(PBKDF2_ALGORITHM);
			secretKey = secretKeyFactory.generateSecret(spec);
		} catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
			e.printStackTrace();
		}
		return secretKey;
	}

	public static void writeSecretKey(SecretKey secretKey) throws IOException {

		Files.write(Paths.get(SECRET_KEY_PATH), secretKey.getEncoded());
		Files.write(Paths.get(SECRET_KEY_BASE64_PATH), Base64.getEncoder().encode(secretKey.getEncoded()));
	}
	
	public static byte[] readSecretKey() throws IOException {
		
		byte[] secretKeyBytes = Files.readAllBytes(Paths.get(SECRET_KEY_PATH));
		return secretKeyBytes;
	}

	public static byte[] generateInitVector() {

		SecureRandom random = new SecureRandom();
		byte[] initVectorBytes = new byte[INIT_VECTOR_LENGTH];
		random.nextBytes(initVectorBytes);

		return initVectorBytes;
	}
	
	public static void writeInitVector(byte[] initVectorBytes) throws IOException {
		
		Files.write(Paths.get(INIT_VECTOR_PATH), initVectorBytes);
		Files.write(Paths.get(INIT_VECTOR_BASE64_PATH), Base64.getEncoder().encode(initVectorBytes));
	}
	
	public static byte[] readInitVector() throws IOException {
		
		byte[] initVectorBytes = Files.readAllBytes(Paths.get(INIT_VECTOR_PATH));
		return initVectorBytes;
	}

}
